from turtle import forward
import paddlenlp as ppnlp
import paddle.nn as pann
import torch.nn as tonn
from transformers import BertModel, ErnieModel

class ernie_classifier(pann.Layer):
    def __init__(self, seg_class, poly_class, rhythm_class, model_name, model_hidden_size):
        super(ernie_classifier, self).__init__()
        self.ernie = ppnlp.transformers.ErnieModel.from_pretrained(model_name)
        self.linear_seg = pann.Linear(model_hidden_size, seg_class)
        self.linear_poly = pann.Linear(model_hidden_size, poly_class)
        self.rhythm_1 = pann.Linear(model_hidden_size, rhythm_class)
        self.rhythm_2 = pann.Linear(model_hidden_size, rhythm_class)
        self.rhythm_3 = pann.Linear(model_hidden_size, rhythm_class)

    def forward(self, input_ids, attention_mask, batch_ids, poly_ids):
        outputs = self.ernie(input_ids, attention_mask = attention_mask)
        seg = self.linear_seg(outputs[0])
        polyphone = self.linear_poly(outputs[0][batch_ids, poly_ids])
        rhythm_1 = self.rhythm_1(outputs[0])
        rhythm_2 = self.rhythm_2(outputs[0])
        rhythm_3 = self.rhythm_3(outputs[0])

        return polyphone, seg, rhythm_1, rhythm_2, rhythm_3



class pt_ernie_classifier(tonn.Module):
    def __init__(self, seg_class, poly_class, rhythm_class, model_name, model_hidden_size):
        super(pt_ernie_classifier, self).__init__()
        self.bert = ErnieModel.from_pretrained(model_name)
        self.linear_seg = tonn.Linear(model_hidden_size, seg_class)
        self.linear_poly = tonn.Linear(model_hidden_size, poly_class)
        self.rhythm_1 = tonn.Linear(model_hidden_size, rhythm_class)
        self.rhythm_2 = tonn.Linear(model_hidden_size, rhythm_class)
        self.rhythm_3 = tonn.Linear(model_hidden_size, rhythm_class)

    def forward(self, input_ids, attention_mask, batch_ids, poly_ids):
        outputs = self.bert(input_ids, attention_mask = attention_mask)
        seg = self.linear_seg(outputs[0])
        polyphone = self.linear_poly(outputs[0][batch_ids, poly_ids])
        rhythm_1 = self.rhythm_1(outputs[0])
        rhythm_2 = self.rhythm_2(outputs[0])
        rhythm_3 = self.rhythm_3(outputs[0])

        return polyphone, seg, rhythm_1, rhythm_2, rhythm_3


        
class pt_bert_classifier(tonn.Module):
    def __init__(self, seg_class, poly_class, rhythm_class, model_name, model_hidden_size):
        super(pt_bert_classifier, self).__init__()
        self.bert = BertModel.from_pretrained(model_name)
        self.linear_seg = tonn.Linear(model_hidden_size, seg_class)
        self.linear_poly = tonn.Linear(model_hidden_size, poly_class)
        self.rhythm_1 = tonn.Linear(model_hidden_size, rhythm_class)
        self.rhythm_2 = tonn.Linear(model_hidden_size, rhythm_class)
        self.rhythm_3 = tonn.Linear(model_hidden_size, rhythm_class)

    def forward(self, input_ids, attention_mask, batch_ids, poly_ids):
        outputs = self.bert(input_ids, attention_mask = attention_mask)
        seg = self.linear_seg(outputs[0])
        polyphone = self.linear_poly(outputs[0][batch_ids, poly_ids])
        rhythm_1 = self.rhythm_1(outputs[0])
        rhythm_2 = self.rhythm_2(outputs[0])
        rhythm_3 = self.rhythm_3(outputs[0])

        return polyphone, seg, rhythm_1, rhythm_2, rhythm_3


class pt_bert_attention_classifier(tonn.Module):
    def __init__(self, seg_class, poly_class, rhythm_class, model_name, model_hidden_size):
        super(pt_bert_attention_classifier, self).__init__()
        self.bert = BertModel.from_pretrained(model_name)
        self.attn_seg = my_multiheadattention(model_hidden_size, 12)
        self.attn_poly = my_multiheadattention(model_hidden_size, 12)
        self.attn_rhythm = my_multiheadattention(model_hidden_size, 12)
        self.linear_seg = tonn.Linear(model_hidden_size, seg_class)
        self.linear_poly = tonn.Linear(model_hidden_size, poly_class)
        self.rhythm_1 = tonn.Linear(model_hidden_size, rhythm_class)
        self.rhythm_2 = tonn.Linear(model_hidden_size, rhythm_class)
        self.rhythm_3 = tonn.Linear(model_hidden_size, rhythm_class)

    def forward(self, input_ids, attention_mask, batch_ids, poly_ids):
        outputs = self.bert(input_ids, attention_mask = attention_mask)[0]
        seg = self.linear_seg(self.attn_seg(outputs))
        polyphone = self.linear_poly(self.attn_poly(outputs)[batch_ids, poly_ids])
        rhythm_1 = self.rhythm_1(self.attn_rhythm(outputs))
        rhythm_2 = self.rhythm_2(self.attn_rhythm(outputs))
        rhythm_3 = self.rhythm_3(self.attn_rhythm(outputs))

        return polyphone, seg, rhythm_1, rhythm_2, rhythm_3

class my_multiheadattention(tonn.Module):
    def __init__(self, input_size, head_nums):
        super(my_multiheadattention, self).__init__()
        self.qweight = tonn.Linear(input_size, input_size)
        self.kweight = tonn.Linear(input_size, input_size)
        self.vweight = tonn.Linear(input_size, input_size)
        self.attn = tonn.MultiheadAttention(input_size, head_nums, batch_first=True)

    def forward(self, input):
        q = self.qweight(input)
        k = self.kweight(input)
        v = self.kweight(input)
        output = self.attn(q, k, v, need_weights=False)
        return output[0]